import numpy as np

def find_topk(a, k, axis=-1, largest=True, sorted=True):
    if axis is None:
        axis_size = a.size
    else:
        axis_size = a.shape[axis]
    assert 1 <= k <= axis_size

    a = np.asanyarray(a)
    if largest:
        index_array = np.argpartition(a, axis_size-k, axis=axis)
        topk_indices = np.take(index_array, -np.arange(k)-1, axis=axis)
    else:
        index_array = np.argpartition(a, k-1, axis=axis)
        topk_indices = np.take(index_array, np.arange(k), axis=axis)
    topk_values = np.take_along_axis(a, topk_indices, axis=axis)
    if sorted:
        sorted_indices_in_topk = np.argsort(topk_values, axis=axis)
        if largest:
            sorted_indices_in_topk = np.flip(sorted_indices_in_topk, axis=axis)
        sorted_topk_values = np.take_along_axis(
            topk_values, sorted_indices_in_topk, axis=axis)
        sorted_topk_indices = np.take_along_axis(
            topk_indices, sorted_indices_in_topk, axis=axis)
        return sorted_topk_values, sorted_topk_indices
    return topk_values, topk_indices


# # find_topk(a, k, axis=-1,
# arr = np.array([[1,2,3],[3,2,1]])
# topk_values,topk_indices = find_topk(arr,1,axis=1)
# print(topk_values)
# print(topk_indices)

# bool_arr=  np.zeros((2,3),dtype=bool)

# bool_arr[np.arange(2),topk_indices] = True
# print(1-bool_arr)
arr1 = np.ones((1,2,3),dtype = np.float32)
arr2 = np.zeros((10,2,3),dtype = np.int64)
arr1[0,0,0] = -10
print(np.minimum(arr1,arr2))